{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6445f581",
   "metadata": {},
   "source": [
    "# `Variable` and `BrainPyObject`\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_math/variables.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_math/variables.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "348b02c2",
   "metadata": {},
   "source": [
    "@[Chaoming Wang](https://github.com/chaoming0625)\n",
    "@[Xiaoyu Chen](mailto:c-xy17@tsinghua.org.cn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e72cc93b",
   "metadata": {},
   "source": [
    "In BrainPy, the JIT compilation for class objects relies on `brainpy.math.Variable`. Moreover, ``brainpy.math.BrainPyObject`` is the container for wrapping ``Variable`` in a class object.\n",
    "\n",
    "In this section, we are going to understand:\n",
    "\n",
    "- What is a ``Variable``?\n",
    "- What is a ``BrainPyObject``?\n",
    "- How to update a ``Variable``?"
   ]
  },
  {
   "cell_type": "code",
   "id": "7188b466",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.256740Z",
     "start_time": "2025-10-06T04:03:20.037173Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "bm.set_platform('cpu')\n",
    "bp.__version__"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.0.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "id": "53b1704b",
   "metadata": {},
   "source": [
    "## ``brainpy.math.Variable``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f7dc2b",
   "metadata": {},
   "source": [
    "``brainpy.math.Variable`` is a pointer referring to a JAX Array. It stores an array as its value. The data in a Variable can be changed during our object-oriented JIT compilation. **If an array is labeled as a Variable, it means that it is a dynamical variable that changes over time.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecb246c1",
   "metadata": {},
   "source": [
    "Arrays that are not marked as Variables will be JIT compiled as static data. Modifications of these arrays will be invalid or cause an error."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c80bdb1",
   "metadata": {},
   "source": [
    "- **Creating a Variable**\n",
    "\n",
    "Passing a tensor into the ``brainpy.math.Variable`` creates a Variable. For example:"
   ]
  },
  {
   "cell_type": "code",
   "id": "9bdceead",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.515511Z",
     "start_time": "2025-10-06T04:03:24.285432Z"
    }
   },
   "source": [
    "b1 = bm.random.random(5)\n",
    "b1"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.7180916 , 0.5760685 , 0.6447711 , 0.92505515, 0.74312174],      dtype=float32)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "id": "d9d16723",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.526704Z",
     "start_time": "2025-10-06T04:03:24.518256Z"
    }
   },
   "source": [
    "b2 = bm.Variable(b1)\n",
    "b2"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Variable(\n",
       "  value=ShapedArray(float32[5]),\n",
       "  _batch_axis=None,\n",
       "  axis_names=None\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "id": "214010c1",
   "metadata": {},
   "source": [
    "- **Accessing the value in a Variable**\n",
    "\n",
    "The data in a Variable can be obtained through ``.value``."
   ]
  },
  {
   "cell_type": "code",
   "id": "a7c53a9a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.544525Z",
     "start_time": "2025-10-06T04:03:24.538826Z"
    }
   },
   "source": [
    "b2.value"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.7180916 , 0.5760685 , 0.6447711 , 0.92505515, 0.74312174],      dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "1c6621b7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.640359Z",
     "start_time": "2025-10-06T04:03:24.602338Z"
    }
   },
   "source": [
    "(b2.value == b1).all()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(True, dtype=bool)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "id": "6b5281a9",
   "metadata": {},
   "source": [
    "- **Supported operations on Variables**\n",
    "\n",
    "Variables support almost all the operations for a JAX array."
   ]
  },
  {
   "cell_type": "code",
   "id": "0824d649",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.707512Z",
     "start_time": "2025-10-06T04:03:24.677241Z"
    }
   },
   "source": [
    "b2 + 1."
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1.7180916, 1.5760685, 1.6447711, 1.9250551, 1.7431217], dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "628fbecc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.742477Z",
     "start_time": "2025-10-06T04:03:24.721153Z"
    }
   },
   "source": [
    "b2 ** 2"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.5156556 , 0.33185494, 0.41572976, 0.855727  , 0.55222994],      dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "id": "4bb90bb0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:24.788958Z",
     "start_time": "2025-10-06T04:03:24.768146Z"
    }
   },
   "source": [
    "bm.floor(b2)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([0., 0., 0., 0., 0.]), dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "id": "f3a42f7a",
   "metadata": {},
   "source": [
    "## ``brainpy.math.BrainPyObject``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d43be23a",
   "metadata": {},
   "source": [
    "``BrainPyObject`` can be viewed as a base class for wrapping all `Variable`s in a class.\n",
    "\n",
    "By using ``.vars()`` function, any ``Variable`` defined in this class can be easily pulled out."
   ]
  },
  {
   "cell_type": "code",
   "id": "3f08bc72",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.175061Z",
     "start_time": "2025-10-06T04:03:24.791977Z"
    }
   },
   "source": [
    "# for example\n",
    "\n",
    "hh = bp.neurons.HH(1)\n",
    "\n",
    "hh.vars().keys()  # a HH model has 6 variables"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['HH0.V', 'HH0.m', 'HH0.h', 'HH0.n', 'HH0.spike', 'HH0.input'])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "id": "f4432226",
   "metadata": {},
   "source": [
    "## Subtypes of ``Variable``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aac1dc5",
   "metadata": {},
   "source": [
    "### Customizing ``Variable`` types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "327c22d5",
   "metadata": {},
   "source": [
    "Sometimes, different variables can be served as the different roles in a computation. To distinguish them from each other, we can subclass ``brainpy.math.Variable`` to define any variable class you want."
   ]
  },
  {
   "cell_type": "code",
   "id": "cb7aeb05",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.198308Z",
     "start_time": "2025-10-06T04:03:25.193538Z"
    }
   },
   "source": [
    "class YourVar(bm.Variable):\n",
    "    pass"
   ],
   "outputs": [],
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "id": "09c2b560",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.207116Z",
     "start_time": "2025-10-06T04:03:25.202759Z"
    }
   },
   "source": [
    "class YourModel(bm.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.a = bm.Variable(bm.ones(1))\n",
    "        self.b = YourVar(bm.zeros(1))"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "id": "b5de262c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.219322Z",
     "start_time": "2025-10-06T04:03:25.214172Z"
    }
   },
   "source": [
    "model = YourModel()\n",
    "model.vars().subset(YourVar)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'YourModel0.b': YourVar(\n",
       "   value=ShapedArray(float32[1]),\n",
       "   _batch_axis=None,\n",
       "   axis_names=None\n",
       " )}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "id": "5c492396",
   "metadata": {},
   "source": [
    "\n",
    "In BrainPy, we have provided many ``Variable`` subtypes, including ``brainpy.math.TrainVar``, ``brainpy.math.Parameter``, and ``brainpy.math.RandomState``."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad677bf0",
   "metadata": {},
   "source": [
    "### 1. TrainVar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5504c217",
   "metadata": {},
   "source": [
    "``brainpy.math.TrainVar`` is a trainable variable and a subclass of ``brainpy.math.Variable``. Usually, the trainable variables are meant to require their gradients and compute the corresponding update values. However, users can also use TrainVar for other purposes. "
   ]
  },
  {
   "cell_type": "code",
   "id": "f8357f81",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.324571Z",
     "start_time": "2025-10-06T04:03:25.231365Z"
    }
   },
   "source": [
    "b = bm.random.rand(4)\n",
    "\n",
    "b"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.7180635 , 0.65054834, 0.819826  , 0.68358934], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "id": "21f05b09",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.339692Z",
     "start_time": "2025-10-06T04:03:25.335579Z"
    }
   },
   "source": [
    "bm.TrainVar(b)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainVar(\n",
       "  value=ShapedArray(float32[4]),\n",
       "  _batch_axis=None,\n",
       "  axis_names=None\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "id": "e8284d53",
   "metadata": {},
   "source": [
    "### 2. Parameter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96aa1cf9",
   "metadata": {},
   "source": [
    "``brainpy.math.Parameter`` is to label a dynamically changed parameter. It is also a subclass of ``brainpy.math.Variable``. The advantage of using Parameter rather than Variable is that it can be easily retrieved by the ``Collector.subset`` method."
   ]
  },
  {
   "cell_type": "code",
   "id": "79105af2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.624351Z",
     "start_time": "2025-10-06T04:03:25.364262Z"
    }
   },
   "source": [
    "b = bm.random.rand(1)\n",
    "\n",
    "b"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.892637], dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 15
  },
  {
   "cell_type": "code",
   "id": "773edf8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.633908Z",
     "start_time": "2025-10-06T04:03:25.630211Z"
    }
   },
   "source": [
    "bm.Parameter(b)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter(\n",
       "  value=ShapedArray(float32[1]),\n",
       "  _batch_axis=None,\n",
       "  axis_names=None\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "id": "afd5dfaa",
   "metadata": {},
   "source": [
    "### 3. RandomState"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba9c30c7",
   "metadata": {},
   "source": [
    "``brainpy.math.random.RandomState`` is also a subclass of ``brainpy.math.Variable``. RandomState must store the dynamically changed **key** information (see [JAX random number designs](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers)). Every time after a RandomState performs a random sampling, the \"key\" will change. Therefore, it is worthy to label a RandomState as the Variable. "
   ]
  },
  {
   "cell_type": "code",
   "id": "e2ce1816",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.687429Z",
     "start_time": "2025-10-06T04:03:25.646128Z"
    }
   },
   "source": [
    "state = bm.random.RandomState(1234)\n",
    "\n",
    "state"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomState([   0 1234])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 17
  },
  {
   "cell_type": "code",
   "id": "b3360505",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.710907Z",
     "start_time": "2025-10-06T04:03:25.704798Z"
    }
   },
   "source": [
    "# perform a \"random\" sampling \n",
    "state.random(1)\n",
    "\n",
    "state  # the value changed"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomState([1264997412 2518116175])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 18
  },
  {
   "cell_type": "code",
   "id": "27dfae54",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.735874Z",
     "start_time": "2025-10-06T04:03:25.731390Z"
    }
   },
   "source": [
    "# perform a \"sample\" sampling \n",
    "state.sample(1)\n",
    "\n",
    "state  # the value changed too"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomState([3834138653 3475645137])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 19
  },
  {
   "cell_type": "markdown",
   "id": "b5bcef7a",
   "metadata": {},
   "source": [
    "Every instance of RandomState can create a new seed from the current seed with ``.split_key()``. "
   ]
  },
  {
   "cell_type": "code",
   "id": "ac30eb3d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:25.761319Z",
     "start_time": "2025-10-06T04:03:25.755197Z"
    }
   },
   "source": [
    "state.split_key()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1364783093, 4258707102], dtype=uint32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 20
  },
  {
   "cell_type": "markdown",
   "id": "f9f8f0fb",
   "metadata": {},
   "source": [
    "It can also create multiple seeds from the current seed with ``.split_keys(n)``. This is used internally by [pmap](../apis/auto/math/generated/brainpy.math.parallels.pmap.rst) and [vmap](../apis/auto/math/generated/brainpy.math.parallels.vmap.rst) to ensure that random numbers are different in parallel threads. "
   ]
  },
  {
   "cell_type": "code",
   "id": "fd164f9e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:03:26.632674Z",
     "start_time": "2025-10-06T04:03:25.776986Z"
    }
   },
   "source": [
    "state.split_keys(2)"
   ],
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'RandomState' object has no attribute 'split_keys'",
     "output_type": "error",
     "traceback": [
      "\u001B[31m---------------------------------------------------------------------------\u001B[39m",
      "\u001B[31mAttributeError\u001B[39m                            Traceback (most recent call last)",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[21]\u001B[39m\u001B[32m, line 1\u001B[39m\n\u001B[32m----> \u001B[39m\u001B[32m1\u001B[39m state.split_keys(\u001B[32m2\u001B[39m)\n",
      "\u001B[31mAttributeError\u001B[39m: 'RandomState' object has no attribute 'split_keys'"
     ]
    }
   ],
   "execution_count": 21
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "32b018e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[3244149147, 2659778815],\n",
       "             [2548793527, 3057026599],\n",
       "             [ 874320145, 4142002431],\n",
       "             [3368470122, 3462971882],\n",
       "             [1756854521, 1662729797]], dtype=uint32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state.split_keys(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bd9149a",
   "metadata": {},
   "source": [
    "There is a default RandomState in ``brainpy.math.random`` module: `DEFAULT`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4f13cfae",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomState(key=([1573595401,  117587871], dtype=uint32))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.random.DEFAULT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75b36c67",
   "metadata": {},
   "source": [
    "The inherent random methods like ``randint()``, ``rand()``, ``shuffle()``, etc. are using this DEFAULT state. If you try to change the default RandomState, please use ``seed()`` method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9c93bdb6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RandomState(key=([     0, 654321], dtype=uint32))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bm.random.seed(654321)\n",
    "\n",
    "bm.random.DEFAULT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10384b23",
   "metadata": {},
   "source": [
    "## How to update ``Variable``?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81cf35f3",
   "metadata": {},
   "source": [
    "In BrainPy, the transformations (like JIT) usually need to update variables or arrays **in-place**. In-place updating does not change the reference pointing to the variable while changing the data stored in the variable."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6b44bda",
   "metadata": {},
   "source": [
    "For example, here we have a variable ``a``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "2c9da6cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = bm.Variable(bm.zeros(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1030c44",
   "metadata": {},
   "source": [
    "The ids of the variable and the data stored in the variable are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "80cce760",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id(a)       =  2151307205504\n",
      "id(a.value) =  2151397524800\n"
     ]
    }
   ],
   "source": [
    "id_of_a = id(a)\n",
    "id_of_data = id(a.value)\n",
    "\n",
    "assert id_of_a != id_of_data\n",
    "\n",
    "print('id(a)       = ', id_of_a)\n",
    "print('id(a.value) = ', id_of_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "217566cb",
   "metadata": {},
   "source": [
    "In-place update (here we use `[:]`) does not change the pointer refered to the variable but changes its data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "01a8e078",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id(a)       =  2151307205504\n",
      "id(a.value) =  2151399078480\n"
     ]
    }
   ],
   "source": [
    "a[:] = 1.\n",
    "\n",
    "print('id(a)       = ', id(a))\n",
    "print('id(a.value) = ', id(a.value))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "29e1c7ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(id(a) == id_of_a)          = True\n",
      "(id(a.value) == id_of_data) = False\n"
     ]
    }
   ],
   "source": [
    "print('(id(a) == id_of_a)          =', id(a) == id_of_a)\n",
    "print('(id(a.value) == id_of_data) =', id(a.value) == id_of_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9d62d23",
   "metadata": {},
   "source": [
    "However, once you do not use in-place operators to assign data, the id that the variable ``a`` refers to will change. This will cause serious errors when using transformations in BrainPy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f20fbb6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "id(a) =  2151305428432\n",
      "(id(a) == id_of_a) = False\n"
     ]
    }
   ],
   "source": [
    "a = 10.\n",
    "\n",
    "print('id(a) = ', id(a))\n",
    "print('(id(a) == id_of_a) =', id(a) == id_of_a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7076ea7",
   "metadata": {},
   "source": [
    "```note\n",
    "The following in-place operators are not limited to ``brainpy.math.Variable`` and its subclasses. They can also apply to ``brainpy.math.Array``. \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44d5bd7",
   "metadata": {},
   "source": [
    "Here, we list several commonly used in-place operators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "00821ab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "v = bm.Variable(bm.arange(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "3c751c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_id = id(v)\n",
    "\n",
    "def check_no_change(new_v):\n",
    "    assert id(new_v) == old_id, 'Variable has been changed.'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d413c648",
   "metadata": {},
   "source": [
    "### 1. Indexing and slicing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93573b82",
   "metadata": {},
   "source": [
    "Indexing and slicing are the two most commonly used operators. The details of indexing and slicing are in [Array Objects Indexing](https://numpy.org/doc/stable/reference/arrays.indexing.html). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3767b90",
   "metadata": {},
   "source": [
    "Indexing: ``v[i] = a`` or  ``v[(1, 3)] = c`` (index multiple values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "87ed7018",
   "metadata": {},
   "outputs": [],
   "source": [
    "v[0] = 1\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fd5ff2f",
   "metadata": {},
   "source": [
    "Slicing: ``v[i:j] = b``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "bbadb60b",
   "metadata": {},
   "outputs": [],
   "source": [
    "v[1: 2] = 1\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750f5203",
   "metadata": {},
   "source": [
    "Slicing all values: ``v[:] = d``, ``v[...] = e``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "4517b203",
   "metadata": {},
   "outputs": [],
   "source": [
    "v[:] = 0\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "dcb6f8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "v[...] = bm.arange(10)\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "076eb1a4",
   "metadata": {},
   "source": [
    "### 2. Augmented assignment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e00a66e",
   "metadata": {},
   "source": [
    "All augmented assignment are in-place operations, which include \n",
    "  - add: ``+=``\n",
    "  - subtract: ``-=``\n",
    "  - divide: ``/=``\n",
    "  - multiply: ``*=``\n",
    "  - floor divide: ``//=``\n",
    "  - modulo: ``%=``\n",
    "  - power: ``**=``\n",
    "  - and: ``&=``\n",
    "  - or: ``|=``\n",
    "  - xor: ``^=``\n",
    "  - left shift: ``<<=``\n",
    "  - right shift: ``>>=``"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "48eea0fa",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "v += 1\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "122eafc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "v *= 2\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "1ff5afc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "v |= bm.random.randint(0, 2, 10)\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "e1625cd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "v **= 2\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "8a46a43c",
   "metadata": {},
   "outputs": [],
   "source": [
    "v >>= 2\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d4ed316",
   "metadata": {},
   "source": [
    "### 3. ``.value`` assignment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34636eb2",
   "metadata": {},
   "source": [
    "Another way to in-place update a variable is to assign new data to ``.value``. This operation is very **safe**, because it will check whether the type and shape of the new data are consistent with the current ones. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "2f81a257",
   "metadata": {},
   "outputs": [],
   "source": [
    "v.value = bm.arange(10)\n",
    "\n",
    "check_no_change(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "19611ce1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'brainpy.errors.MathError'> The shape of the original data is (10,), while we got () with batch_axis=None.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    v.value = bm.asarray(1.)\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "c7911157",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'brainpy.errors.MathError'> The dtype of the original data is int32, while we got float32.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    v.value = bm.random.random(10)\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efd1dcbd",
   "metadata": {},
   "source": [
    "### 4. ``.update()`` method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29630aaa",
   "metadata": {},
   "source": [
    "Actually, the ``.value`` assignment is the same operation as the ``.update()`` method. Users who want a safe assignment can choose this method too. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "d861440c",
   "metadata": {},
   "outputs": [],
   "source": [
    "v.update(bm.random.randint(0, 20, size=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "247f081b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'brainpy.errors.MathError'> The shape of the original data is (10,), while we got () with batch_axis=None.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    v.update(bm.asarray(1.))\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "9ae0ce26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'brainpy.errors.MathError'> The dtype of the original data is int32, while we got float32.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    v.update(bm.random.random(10))\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
